import os
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
from tqdm import tqdm
import logging
import torchvision
import pandas as pd

# 日志设置
from util.log_util import get_logger

get_logger(logging.INFO)

# 设备设置：有GPU用GPU，没有就用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")

# Patch图像文件夹路径（你根据实际路径修改）
img_dir = r"D:/zm/PGGN_Pic"  # 修改为预处理后的patch路径
save_dir = r"D:/zm/PGGN_Pic_feature_file"  # 特征npy文件保存目录
csv_dir = r"D:/zm/PGGN_csv"  # 特征csv文件保存目录
# 如果指定的文件夹不存在，就自动创建它；如果已经存在，就跳过，不报错。
os.makedirs(save_dir, exist_ok=True)
# 如果指定的文件夹不存在，就自动创建它；如果已经存在，就跳过，不报错。
os.makedirs(csv_dir, exist_ok=True)


# 自定义数据集
class PatchDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.files = [f for f in os.listdir(img_dir) if f.endswith('.png') or f.endswith('.jpg')]
        self.transform = transform

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img_name = self.files[idx]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, img_name


# 图像预处理：与ResNet预训练模型对齐
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

# 数据加载器
dataset = PatchDataset(img_dir, transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=False)

# 加载ResNet50模型（去掉分类层，只保留特征提取部分）
resnet50 = models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
feature_extractor = torch.nn.Sequential(*list(resnet50.children())[:-1])  # 去掉最后的fc层
feature_extractor = feature_extractor.to(device)
feature_extractor.eval()

# 提取特征并保存为CSV和NPY
all_features = []
all_names = []

# 提取并保存特征
with torch.no_grad():
    for images, names in tqdm(dataloader):
        images = images.to(device)
        features = feature_extractor(images)  # 输出 shape: [batch_size, 2048, 1, 1]
        features = features.squeeze(-1).squeeze(-1).cpu().numpy()  # shape: [batch_size, 2048]

        for i, name in enumerate(names):
            np.save(os.path.join(save_dir, name.replace(".png", ".npy").replace(".jpg", ".npy")), features[i])

        # 收集所有特征用于CSV
        all_features.extend(features)
        all_names.extend(names)

# 生成CSV文件
df = pd.DataFrame(all_features, index=all_names)
df.columns = [f"feature_{i}" for i in range(df.shape[1])]
csv_path = os.path.join(csv_dir, "features.csv")
df.to_csv(csv_path, index_label="image_name")
logging.info("特征已保存为CSV文件: %s", csv_path)
